{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Install the Package\n",
    "Here we're installing it directly from GitHub while it's in development."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install 'vanna[flask,anthropic]'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Download a Sample Database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import httpx\n",
    "\n",
    "with open(\"Chinook.sqlite\", \"wb\") as f:\n",
    "    with httpx.stream(\"GET\", \"https://vanna.ai/Chinook.sqlite\") as response:\n",
    "        for chunk in response.iter_bytes():\n",
    "            f.write(chunk)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from vanna import Agent, AgentConfig\n",
    "from vanna.servers.fastapi import VannaFastAPIServer\n",
    "from vanna.core.registry import ToolRegistry\n",
    "from vanna.core.user import UserResolver, User, RequestContext\n",
    "from vanna.integrations.anthropic import AnthropicLlmService\n",
    "from vanna.tools import RunSqlTool, VisualizeDataTool\n",
    "from vanna.integrations.sqlite import SqliteRunner\n",
    "from vanna.tools.agent_memory import SaveQuestionToolArgsTool, SearchSavedCorrectToolUsesTool\n",
    "from vanna.integrations.local.agent_memory import DemoAgentMemory\n",
    "from vanna.capabilities.sql_runner import RunSqlToolArgs\n",
    "from vanna.tools.visualize_data import VisualizeDataArgs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define your User Authentication\n",
    "Here we're going to say that if you're logged in as `admin@example.com` then you're in the `admin` group, otherwise you're in the `user` group"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleUserResolver(UserResolver):\n",
    "    async def resolve_user(self, request_context: RequestContext) -> User:\n",
    "        # In production, validate cookies/JWTs here\n",
    "        user_email = request_context.get_cookie('vanna_email')\n",
    "        if not user_email:\n",
    "            raise ValueError(\"Missing 'vanna_email' cookie for user identification\")\n",
    "        \n",
    "        print(f\"Resolving user for email: {user_email}\")\n",
    "\n",
    "        if user_email == \"admin@example.com\":\n",
    "            return User(id=\"admin1\", email=user_email, group_memberships=['admin'])\n",
    "        \n",
    "        return User(id=\"user1\", email=user_email, group_memberships=['user'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define the Tools and Access Control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = ToolRegistry()\n",
    "tools.register_local_tool(RunSqlTool(sql_runner=SqliteRunner(database_path=\"./Chinook.sqlite\")), access_groups=['admin', 'user'])\n",
    "tools.register_local_tool(VisualizeDataTool(), access_groups=['admin', 'user'])\n",
    "agent_memory = DemoAgentMemory(max_items=1000)\n",
    "tools.register_local_tool(SaveQuestionToolArgsTool(), access_groups=['admin'])\n",
    "tools.register_local_tool(SearchSavedCorrectToolUsesTool(), access_groups=['admin', 'user'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up LLM\n",
    "llm = AnthropicLlmService(model=\"claude-sonnet-4-5\", api_key=\"sk-ant-...\")\n",
    "\n",
    "# Create agent with your options\n",
    "agent = Agent(\n",
    "    llm_service=llm,\n",
    "    tool_registry=tools,\n",
    "    user_resolver=SimpleUserResolver(),\n",
    "    config=AgentConfig(),\n",
    "    agent_memory=agent_memory\n",
    ")\n",
    "\n",
    "# 4. Create and run server\n",
    "server = VannaFastAPIServer(agent)\n",
    "server.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
